package com.sunday.common.mq.rocket.brave.propagation;

import brave.Span;
import brave.Tracer;
import org.apache.rocketmq.client.trace.TraceConstants;

/**
 * pull brave threadLocal
 *
 * @author zsy
 * @since 2023/8/17
 */
public class TracingRocketPostProcessor {

    static final ThreadLocal<SpansAndScope> local = new ThreadLocal<>();

    /**
     * 释放
     */
    protected static void release(SpansAndScope spansAndScope, Throwable error) {
        if (spansAndScope == null) {
            return;
        }
        try {
            if (spansAndScope.spans() != null && spansAndScope.spans().size() > 0) {
                for (Span span : spansAndScope.spans()) {
                    if (!span.isNoop()) {
                        if (error != null) {
                            span.error(error);
                            maybeTag(span, TraceConstants.ROCKETMQ_SUCCESS, Boolean.FALSE.toString());
                        } else {
                            maybeTag(span, TraceConstants.ROCKETMQ_SUCCESS, Boolean.TRUE.toString());
                        }
                        span.finish();
                    }
                }
            }
        } finally {
            Tracer.SpanInScope spanInScope = spansAndScope.scope();
            if (spanInScope != null) {
                spanInScope.close();
            }
        }
    }

    static void maybeTag(Span span, String tag, String value) {
        if (value != null) span.tag(tag, value);
    }


    /**
     * 批量
     */
    public static void release(Throwable error) {
        release(local.get(), error);
    }

    /**
     * 设置
     */
    public static void set(SpansAndScope spansAndScope) {
        local.set(spansAndScope);
    }

    /**
     * 结束范围
     */
    public static void close(Throwable error) {
        release(local.get(), error);
        local.remove();
    }

}
